import torch
import torch.nn as nn

device = torch.device('cuda')


class HomoCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(

            make_layers(36, 108),
            make_layers(108, 324),
            # make_layers(216, 384),
            # make_layers(324, 108),

            # make_layers(36, 108),
            # make_layers(108, 216),
            # # make_layers(216, 384),
            # make_layers(216, 432),
            nn.Conv3d(324, 108, kernel_size=3, stride=1, padding=1),
            nn.Conv3d(108, 18, kernel_size=3, stride=1, padding=1),
        )

    def forward(self, X):
        return self.conv(X)


def make_layers(in_channels, out_channels):
    conv = nn.Sequential(
        nn.BatchNorm3d(in_channels),
        nn.Conv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),
        nn.Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
        nn.BatchNorm3d(out_channels),
        nn.LeakyReLU(inplace=True),
        # nn.Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
    )

    return conv
